#!/usr/bin/env python
#Copyright 2009 Erik Tollerud
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
This script runs on fits and fits.gz files and extracts header  or summary 
information from the files.  If no files are given it looks at all fits and 
fits.gz files in the current directory.
"""

import sys
from optparse import OptionParser
from glob import glob

try:
    import pyfits
except ImportError:
    print 'PyFits not found, cannot continue - exiting...'
    sys.exit(1)

op = OptionParser()
op.usage = '%prog [options] [filename(s)]'
op.add_option('-e','--extension',dest='ext',help='Extension of fits file from which to read header.  If not specified, a summary will be shown.',metavar='EXT',default=None)
op.add_option('-r','--record',dest='rec',help='Record names to lookup',metavar='REC')
op.add_option('-c','--coords',dest='coords',help='Decimal coordinate lookup',action='store_true',default=False)
op.add_option('-f','--flat',dest='flat',help='Flat listing for each entry',action='store_true',default=False)

ops,args=op.parse_args()
if len(args) < 1:
    args.append('*.fits')
    args.append('*.fits.gz')
    
fns = []
for a in args:
    fns.extend(glob(a))
    
if ops.ext is None:
    rec = ops.rec
    if rec is not None and ',' in rec:
        print "Can't show multiple records in summary mode"
        sys.exit(2)
    
    dosep = False
    for fn in fns:
        if dosep:
            print '-------------------------------------------','\n'
        if not (fn.endswith('fits') or fn.endswith('fits.gz')):
            print 'File',fn,'is not FITS, skipping...','\n'
        else:
            try:
                f = pyfits.open(fn,memmap=1)
                try:
                    print 'File',fn,'contains',len(f),'HDU'+('s' if len(f)>1 else ''),':'
                    for i,hdu in enumerate(f):
                        print 'HDU #%i, name:%s, type:%s'%(i,hdu.name,hdu.__class__.__name__)
                        print '\tHeader has %i entries'%len(hdu.header.keys())
                        if rec:
                            if hdu.header.has_key(rec):
                                print '\t',rec,'=',hdu.header.get(rec)
                            else:
                                print '\t',rec,'not found in this header'
                        if hdu.data is None:
                            print '\tContains no Data'
                        else:
                            print '\tData shape:',hdu.data.shape
                finally:
                    f.close()
                dosep = True
            except Exception,e:
                print 'Problem summarizing file',fn+':',str(e)
else:
    flat = ops.flat    
    ext = int(ops.ext)
    rec = ops.rec
    if rec is not None:
        if ops.coords:
            print "Can't use coordinate mode and records to lookup, exiting"
            sys.exit(2)
        rec = rec.split(',')
    elif ops.coords:
        from astropysics.coords import EquatorialPosition
        rec = ['RA','DEC']
        ap = EquatorialPosition()
        aplast = None

    for fn in fns:
        if not (fn.endswith('fits') or fn.endswith('fits.gz')):
            print 'File',fn,'is not FITS, skipping...'
        else:
            h = pyfits.getheader(fn,ext)
            if not flat:
                print 'File',fn,'HDU #',ext,':'
            if rec is None:
                print str(h).strip()
            else:
                for r in rec:
                    val = h.get(r)
                        
                    if val is None:
                        if flat:
                            print fn+':',r,'Record does not exist'
                        else:
                            print '\t',r,'Record does not exist'
                    else:
                        if ops.coords:                        
                            if r == 'RA':
                                ap.ra.hms = [float(e) for e in val.split(':')]
                                val = ap.ra.d
                                if aplast == 'DEC':
                                    if flat:
                                        print fn+':','coords',ap.ra.d,ap.dec.d
                                    else:
                                        print '\tcoords',ap.ra.d,ap.dec.d
                                    aplast = None
                                else:
                                    aplast = 'RA'
                            elif r == 'DEC':
                                ap.dec.dms = [float(e) for e in val.split(':')]
                                val = ap.dec.d
                                if aplast == 'RA':
                                    if flat:
                                        print fn+':','coords',ap.ra.d,ap.dec.d
                                    else:
                                        print '\tcoords',ap.ra.d,ap.dec.d
                                    aplast = None
                                else:
                                    aplast = 'DEC'
                            else:
                                raise ValueError("This should be unreachable!")
                        if flat:
                            print fn+':',r,'=',val
                        else:    
                            print '\t',r,'=',val
            if not flat:            
                print '' #newline to split files
    
